#include <opencv2/opencv.hpp>
#include <opencv2/highgui.hpp>
#include <torch/script.h>
#include <torch/torch.h>
#include <ATen/ATen.h>
#include <iostream>
#include <map>

std::map<int, std::string> CHARS = {
    {0, "京"}, {1, "沪"}, {2, "津"}, {3, "渝"}, 
    {4, "冀"}, {5, "晋"}, {6, "蒙"}, {7, "辽"}, 
    {8, "吉"}, {9, "黑"}, {10, "苏"}, {11, "浙"}, 
    {12, "皖"}, {13, "闽"}, {14, "赣"}, {15, "鲁"}, 
    {16, "豫"}, {17, "鄂"}, {18, "湘"}, {19, "粤"}, 
    {20, "桂"}, {21, "琼"}, {22, "川"}, {23, "贵"}, 
    {24, "云"}, {25, "藏"}, {26, "陕"}, {27, "甘"}, 
    {28, "青"}, {29, "宁"}, {30, "新"}, {31, "0"}, 
    {32, "1"}, {33, "2"}, {34, "3"}, {35, "4"}, 
    {36, "5"}, {37, "6"}, {38, "7"}, {39, "8"}, 
    {40, "9"}, {41, "A"}, {42, "B"}, {43, "C"}, 
    {44, "D"}, {45, "E"}, {46, "F"}, {47, "G"}, 
    {48, "H"}, {49, "J"}, {50, "K"}, {51, "L"}, 
    {52, "M"}, {53, "N"}, {54, "P"}, {55, "Q"}, 
    {56, "R"}, {57, "S"}, {58, "T"}, {59, "U"}, 
    {60, "V"}, {61, "W"}, {62, "X"}, {63, "Y"}, 
    {64, "Z"}, {65, "I"}, {66, "O"}, {67, "-"}};


int main(int argc, const char *argv[])
{
    // 加载模型
    torch::jit::script::Module module;
    try
    {
        module = torch::jit::load("/home/sunqiang/PROJECT/vscode/testcmake/Final_LPRNet_model.pt");
    }
    catch (const c10::Error &e)
    {
        std::cerr << "error loading the model\n";
        return -1;
    }

    // 图片预处理
    std::string img_path = "/home/sunqiang/PROJECT/vscode/testcmake/images/皖A1R029.jpg";
    cv::Mat img = cv::imread(img_path);
    cv::resize(img, img, cv::Size(94, 24));
    img.convertTo(img, CV_32F);
    img -= 127.5;
    img *= 0.0078125;
    img = img.t();
    auto img_tensor = torch::from_blob(img.data, {1, img.rows, img.cols, img.channels()}, torch::kFloat32);
    img_tensor = img_tensor.permute({0, 3, 2, 1});

    // 模型预测
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(img_tensor);
    torch::Tensor output = module.forward(inputs).toTensor();

    // tensor格式转为vector
    torch::Tensor output_tensor = output.clone();
    output_tensor = output_tensor.permute({0, 2, 1});
    output_tensor = output_tensor.flatten(0, 2);

    std::vector<float> output_vector(output_tensor.data_ptr<float>(), output_tensor.data_ptr<float>() + output_tensor.numel());

    std::vector<std::vector<int>> preb_labels;
    int BatchSize = output.size(0);

    for (int i=0;i < BatchSize;i++)
    {
        std::vector<int> preb_label;
        for (int j = 0;j < 18;j++){
            int max_index = 0;
            float max_value = 0;
            int base_index = i * 18 * 68 + j * 68; 
            for (int k = 0;k < 68;k++){
                float value = output_vector[base_index + k]; 
                if (value > max_value){
                    max_value = value;
                    max_index = k;
                }
            }
            preb_label.push_back(max_index);
        }
        std::vector<int> no_repeat_blank_label;
        int pre_c = preb_label[0];
        if (pre_c != CHARS.size() - 1){
            no_repeat_blank_label.push_back(pre_c);
        }
        for (auto c : preb_label){
            if (pre_c == c || c == CHARS.size() - 1){
                if (c == CHARS.size() - 1){
                    pre_c = c;
                }
                continue;
            }
            no_repeat_blank_label.push_back(c);
            pre_c = c;
        }
        preb_labels.push_back(no_repeat_blank_label);
    }
    
    for (auto label : preb_labels)
    {
        std::cout << "预测结果：";
        for (auto i : label){
            std::cout << CHARS[i];
        }
        std::cout << std::endl;
    }

    /*
    // tensor作为输入
    std::vector<std::vector<int>> preb_labels;
    output = output.cpu(); // Assuming cpu is a method that moves the tensor to cpu

    for (int i = 0; i < output.size(0); i++)
    {
        auto preb = output.index({i, torch::indexing::Slice(torch::indexing::None), torch::indexing::Slice(torch::indexing::None)}); // Assuming slice is a method that slices the tensor
        std::vector<int> preb_label;
        for (int j = 0; j < preb.size(1); j++)
        {
            preb_label.push_back(preb.index({torch::indexing::Slice(torch::indexing::None), j}).argmax(0).item<int>()); // Assuming argmax is a method that returns the index of the max value
        }
        std::vector<int> no_repeat_blank_label;
        int pre_c = preb_label[0];

        if (pre_c != CHARS.size() - 1)
        {
            no_repeat_blank_label.push_back(pre_c);
        }
        for (auto c : preb_label)
        {
            if (pre_c == c || c == CHARS.size() - 1)
            {
                if (c == CHARS.size() - 1)
                {
                    pre_c = c;
                }
                continue;
            }
            no_repeat_blank_label.push_back(c);
            pre_c = c;
        }
        preb_labels.push_back(no_repeat_blank_label);
    }

    for (auto label : preb_labels)
    {
        std::cout << "预测结果：";
        for (auto i : label)
        {
            std::cout << CHARS[i];
        }
        std::cout << std::endl;
    }
    */

    std::string img_name = img_path.substr(img_path.find_last_of('/') + 1);
    std::string real_label = img_name.substr(0, img_name.find_last_of('.'));
    std::cout << "真实结果：" << real_label << std::endl;
    return 0;
}